from mujoco_py import load_model_from_xml, load_model_from_path, MjSim, MjViewer, MjRenderContextOffscreen
import os
import cv2

BASIC_MODEL_XML = """
<mujoco>
    <asset>
        <texture name="t1" width="33" height="36" type="2d" builtin="flat" />
        <texture name="t2" width="34" height="39" type="2d" builtin="flat" />
        <texture name="t3" width="31" height="37" type="2d" builtin="flat" />
        <texture name="t4" width="38" height="34" type="2d" builtin="flat" />
        <material name="m1" texture="t1" />
        <material name="m2" texture="t2" />
        <material name="m3" texture="t3" />
        <material name="m4" texture="t4" />
    </asset>
    <worldbody>
        <light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
        <camera name="xtion" pos="0 0 2.5" zaxis="0 0 1" />
        <geom name="g1" pos="-0.5 0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m1" />
        <geom name="g2" pos="0.5 0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m2" />
        <geom name="g3" pos="0.5 -0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m3" />
        <geom name="g4" pos="-0.5 -0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m4" />
        <!-- The following boxes will show up in reflections -->
        <geom pos="0.5 0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
        <geom pos="0.5 -0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
        <geom pos="-0.5 -0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
        <geom pos="-0.5 0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
    </worldbody>
</mujoco>
"""

if __name__ == '__main__':
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    viewer = MjViewer(sim)
    offscreen = MjRenderContextOffscreen(sim, 0)
    for i in range(3000):
        if i == 500 or i == 1500 or i == 2000:
            offscreen.render(1920, 1080, 0)
            rgb = offscreen.read_pixels(1920, 1080)[0]
            cv2.imwrite("cam/obs_" + str(i) + ".png", rgb)
            print('c')
        viewer.render()
        sim.step()
